# pip install faiss-cpu  -i https://mirrors.aliyun.com/pypi/simple/  # 如果你的环境支持GPU，也可以安装faiss-gpu
import face_recognition
from tool import getFaces
import cv2
import numpy as np
import faiss
import time


def build_faiss_index(faceNameArr):
    """
    构建Faiss索引。
    :param faceNameArr: 包含所有人脸编码和对应名字的字典列表。
    :return: Faiss索引和对应的姓名列表。
    """
    # 提取所有的人脸编码，并转换为适合Faiss的格式（float32类型的numpy数组）
    face_encodings = [faceName["face"] for faceName in faceNameArr]
    face_encodings_np = np.array(face_encodings).astype("float32")

    # 初始化Faiss索引，这里我们使用L2距离（欧式距离），维度为128（FaceNet输出的维度）
    index = faiss.IndexFlatL2(128)
    index.add(face_encodings_np)

    # 返回索引和对应的姓名列表
    names = [faceName["name"] for faceName in faceNameArr]
    return index, names


cap = cv2.VideoCapture(0)

# 假设这是首次加载所有人脸数据并构建Faiss索引
faceNameArr = getFaces()
index, names = build_faiss_index(faceNameArr)

while cap.isOpened():
    ret, img = cap.read()
    if ret:
        cv2.imshow("capture", img)
        key = cv2.waitKey(100)
        if key == ord("a") or key == ord("A"):
            start_time = time.time()
            unknown_image = img.copy()
            face_encodings_list = face_recognition.face_encodings(unknown_image)
            if not face_encodings_list:
                print(f"警告: 未检测到人脸")
                continue

            unknown_encoding = (
                np.array(face_encodings_list[0]).astype("float32").reshape(1, -1)
            )

            # 使用Faiss进行搜索，返回最接近的结果及其距离
            D, I = index.search(unknown_encoding, k=1)

            # 如果最小的距离小于设定的阈值，则认为找到了匹配
            if (
                D[0][0] < 0.4 * 0.4
            ):  # 注意这里调整了阈值计算方式，因为Faiss默认使用平方欧氏距离
                matched_name = names[I[0][0]]
                print(matched_name)
            else:
                print("未找到匹配")
            end_time = time.time()
            print(f"耗时: {end_time - start_time:.6f} 秒")
            break
cap.release()
cv2.destroyAllWindows()
